/*
 * GDPC format (archive might be embedded in executables, for both PE and ELF it's in the section named "pck"):
 *
 * Starts with a 4 bytes magic "GDPC" followed by the format's version on 4 bytes. Then 3 * 4 bytes engine version follows.
 * If version is 2 or above, flags on 4 bytes comes. Bit 0 of flags supposed to indicate encrypted directory. This is followed
 * by the offset of the files data (at 0x14 for version 1 and at 0x18 for version 2 and above) on 8 bytes. This is important
 * because all file offsets in the records are relative to this (could be different if embedded in an executable). Then 64 bytes
 * reserved, usually all zeros. The number of files is at 0x60 (or 0x54 for version 1) on 4 bytes, followed by the file records,
 * numfiles many (maybe encrypted with AES-256-CFB, but I haven't seen such a pack yet). Each file records goes like:
 *
 * 4 bytes name length + padding, followed by an UTF-8 string (might not be zero terminated if its length is multiple of 4).
 * Then the relative offset of file contents on 8 bytes and file size on 8 bytes. After that on 16 bytes, MD5 checksum of
 * contents. For file version 2 and above a flag on 4 bytes ends the record. Bit 0 means encrypted, and if so, then at the given
 * offset content starts with a header: 16 bytes MD5 checksum (same as in the file table record), on 8 bytes the file content's
 * size again (this time without this header, so 40 bytes less as in the file table record), and then 16 bytes random, which is
 * used as AES initial vector. For version 1 and if flag is zero, then this header is missing and offset points to the first byte
 * of the file content. Filenames are usually prefixed with "res://", and there's a hidden ".godot" directory too in the archive.
 *
 * Finding the encryption key is super easy, same method as with RPG Maker. First, look for an encrypted file with a known magic.
 * (Luckily Godot developers are morrons, so they have added a common header and magic to all assets...) Then try to decode the
 * encrypted file with different keys until you get one that decrypts to that magic. Normally AES-256-CFB uses a 32 bytes long
 * key, so that could take for a while, but we can drastically reduce the possibilities and the required time considering these:
 *  1. the key must be in the executable file. That's just (file size - 32) possibilities instead of 2^256 to start with
 *  2. it's in one of the data sections, by parsing the executable format, we can rule out about 70% - 80% of the file
 *  3. we also know that on 64 bit compilers use an 8 bytes alignment, so we can further cut that number down to 1 / 8th
 *  4. no need to run AES on the entire file, a single call per key to decrypt the very first block (16 bytes) suffice.
 *
 * As a result, Godot encryption can be cracked within a few seconds, no need for injecting DLLs or disassembling code and such.
 *
 * NOTE: if the file table would be encrypted as well, then use the same method, just look for "res://" as magic (first entry's
 * filename starts in the very first AES block too). Haven't implemented because I don't have such a game file to test with.
 *
 */

#include "main.h"

#define GDPCMAGIC "GDPC"

/* PCK format */
typedef struct {
    uint8_t magic[4];
    uint32_t version;
    uint32_t engine_major;
    uint32_t engine_minor;
    uint32_t engine_rev;
    uint32_t flags; /* encrypted file table */
} gdpc_t;

static uint64_t total, current;
static uint8_t detected = 0, haskey = 0, godot_key[32] = { 0 };
static void *exe = NULL;
static section_t exesections[64];

/* file index */
typedef struct {
    char name[NAMELEN];
    uint64_t offs;
    uint64_t size;
    int encrypted;
} index_t;
static index_t *idx = NULL;
static int numidx = 0;

/**
 * Detect encryption key
 */
static void detectenckey(index_t *idxrec, uint8_t *inbuf)
{
    mbedtls_aes_context ctx;
    size_t iv_off = 0;
    uint64_t i, len;
    uint8_t iv[16], magic[32], *buf;
    char *ext, *magic1 = NULL, *magic2 = NULL;
    int j;

    if(!idxrec || !idxrec->name[0] || !inbuf || !exe || idxrec->size < 40) return;
    /* try to guess the magic from extension */
    if((ext = strrchr(idxrec->name, '.'))) {
        if(!strcmp(ext, ".res") || !strcmp(ext, ".oggstr") || !strcmp(ext, ".sample")) { magic1 = "RSRC"; } else
        if(!strcmp(ext, ".stex") || !strcmp(ext, ".ctex")) { magic1 = "GST2"; magic2 = "GDST"; } else
        if(!strcmp(ext, ".remap")) { magic1 = "[rem"; }
    }
    if(magic1) {
        status(DETECTKEY, 0, 100);
        for(j = 0; !quit && exesections[j].name[0] && !haskey; j++)
            if(memcmp(exesections[j].name, "pck", 4)) {
                if((buf = (uint8_t*)malloc(exesections[j].size))) {
                    fileseek(exe, exesections[j].offs);
                    len = fileread(exe, buf, exesections[j].size);
                    /* try out contents of data sections as key */
                    for(i = 0; !quit && i + 32 < len && !haskey; i += 8) {
                        iv_off = 0;
                        memcpy(iv, inbuf + 16 + 8, 16);
                        memset(magic, 0, 16);
                        memcpy(magic + 16, inbuf + 16 + 8 + 16, 16);
                        mbedtls_aes_init(&ctx);
                        mbedtls_aes_setkey_enc(&ctx, buf + i, 256);
                        mbedtls_aes_crypt_cfb128(&ctx, MBEDTLS_AES_DECRYPT, 16, &iv_off, iv, magic + 16, magic);
                        mbedtls_aes_free(&ctx);
                        /* bingo! */
                        if(!memcmp(magic, magic1, 4) || (magic2 && !memcmp(magic, magic2, 4))) {
                            memcpy(godot_key, buf + i, sizeof(godot_key));
                            haskey = 1;
                            break;
                        }
                    }
                    free(buf);
                } else {
                    engine_features |= EF_NOMEM;
                    debug("unable to allocate memory for data section %lx (%lu bytes)", exesections[j].offs, exesections[j].size);
                }
            }
        fileclose(exe); exe = NULL;
        if(haskey)
            debug("encryption key "
                "%02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X "
                "%02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X",
                godot_key[0], godot_key[1], godot_key[2], godot_key[3],
                godot_key[4], godot_key[5], godot_key[6], godot_key[7],
                godot_key[8], godot_key[9], godot_key[10], godot_key[11],
                godot_key[12], godot_key[13], godot_key[14], godot_key[15],
                godot_key[16], godot_key[17], godot_key[18], godot_key[19],
                godot_key[20], godot_key[21], godot_key[22], godot_key[23],
                godot_key[24], godot_key[25], godot_key[26], godot_key[27],
                godot_key[28], godot_key[29], godot_key[30], godot_key[31]);
        else
            debug("failed to retrieve encryption key, press F2 and specify manually");
        status(LOADING, 0, 100);
    }
}

/**
 * Decrypt a buffer
 */
static uint64_t decrypt_file(uint8_t *buf, uint64_t len)
{
    mbedtls_aes_context ctx;
    size_t iv_off = 0;
    uint64_t orig;
    uint8_t iv[16];

    if(!haskey || !buf || len < 16 + 8 + 16) return len;
    len -= 16 + 8 + 16;
    memcpy(&orig, buf + 16, 8);
    if(orig > len) orig = len;
    memcpy(iv, buf + 16 + 8, 16);
    mbedtls_aes_init(&ctx);
    mbedtls_aes_setkey_enc(&ctx, godot_key, 256);
    mbedtls_aes_crypt_cfb128(&ctx, MBEDTLS_AES_DECRYPT, (orig + 15) & ~15, &iv_off, iv, buf + 16 + 8 + 16, buf);
    mbedtls_aes_free(&ctx);
    return orig;
}

/**
 * Free file index
 */
static void freefileindex(void)
{
    if(idx) { free(idx); idx = NULL; }
    numidx = 0;
}

/**
 * Get file index (parse PCK archive)
 */
static int getfileindex(void *f, uint64_t offs, uint64_t size)
{
    uint8_t buf[128 + 65536], tmp[1024], *ptr, *ptr2;
    gdpc_t *gdpc = (gdpc_t*)buf;
    uint64_t dataoffs;
    uint32_t numfiles, i, n;
    index_t idxrec;

    freefileindex();
    if(!f || !size) return 0;
    memset(buf, 0, sizeof(buf));
    fileseek(f, offs);
    fileread(f, buf, sizeof(buf));
    if(memcmp(buf, GDPCMAGIC, 4)) {
        debug("bad Godot magic");
        return 0;
    }
    ptr = buf + (gdpc->version >= 2 ? 0x64 : 0x58);
    memcpy(&numfiles, ptr - 4, 4);
    memcpy(&dataoffs, buf + 24, 8);
    debug("Godot Pack offs %lx size %lu, version %u, engine %u.%u.%u, data offs %lx, numfiles %u", offs, size, gdpc->version,
        gdpc->engine_major, gdpc->engine_minor, gdpc->engine_rev, dataoffs, numfiles);
    if(gdpc->version >= 2 && (gdpc->flags & 1)) {
        debug("encrypted directory not supported");
        return 0;
    }
    for(i = 0; !quit && i < numfiles; i++) {
        memset(&idxrec, 0, sizeof(idxrec));
        memcpy(&n, ptr, 4); ptr += 4;
        if(!memcmp(ptr, "res://", 6)) { ptr += 6; n -= 6; }
        memcpy(idxrec.name, ptr, n > (NAMELEN-1) ? (NAMELEN-1) : n); ptr += n;
        if(idxrec.name[0] == '.') idxrec.name[0] = '_';
        memcpy(&idxrec.offs, ptr, sizeof(idxrec.offs)); ptr += 8; idxrec.offs += dataoffs;
        memcpy(&idxrec.size, ptr, sizeof(idxrec.size)); ptr += 8;
        ptr2 = ptr; /* md5 */ ptr += 16;
        if(gdpc->version >= 2) { memcpy(&idxrec.encrypted, ptr, 4); ptr += 4; } else idxrec.encrypted = 0;
        if(idxrec.encrypted) {
            idxrec.size += 16 + 8 + 16;
            if(!haskey && exe) {
                fileseek(f, idxrec.offs);
                fileread(f, tmp, sizeof(tmp));
                if(!memcmp(tmp, ptr2, 16))
                    detectenckey(&idxrec, tmp);
            }
        }
        if(!(idx = (index_t*)realloc(idx, (numidx + 1) * sizeof(index_t)))) {
            debug("unable to allocate memory, too many table entries %u", numidx);
            numidx = 0; return 0;
        }
        memcpy(&idx[numidx++], &idxrec, sizeof(index_t));
        if(verbose > 1) debug("file in archive '%s' (%lu bytes)", idxrec.name, idxrec.size);
    }
    return numidx;
}

/**
 * Callback function for find() in godot_detect()
 */
int godot_detectcb(char *path)
{
    uint8_t tmp[32];
    uint64_t size = 0;
    void *f = fileopen(path, &size);
    int i;

    if(f) {
        fileread(f, tmp, sizeof(tmp));
        fileclose(f);
        if(size > 32 && !memcmp(tmp, GDPCMAGIC, 4)) {
            detected = 1;
            return 0;
        }
        if(elfpe_getsections(path, 0))
            for(i = 0; sections[i].name[0]; i++)
                if(!memcmp(sections[i].name, "pck", 4)) {
                    detected = 1;
                    return 0;
                }
    }
    return 1;
}

/**
 * Detect if this is a Godot game
 */
int godot_detect(char *path, uint8_t *buf, uint64_t size)
{
    char tmp[PATH_MAX + FILENAME_MAX + 1];

    if(!path || !*path) return 0;
    if(buf && size > 32 && !memcmp(buf, GDPCMAGIC, 4)) return 1;

    strcpy(tmp, path);
    detected = 0;
    find(tmp, 0, godot_detectcb);
    return detected;
}

/**
 * Callback function for find() in godot_load()
 */
int godot_findexecb(char *path)
{
    uint64_t size = 0;
    int i;

    if(!exe && elfpe_getsections(path, 0))
        for(i = 0; sections[i].name[0]; i++)
            if(!memcmp(sections[i].name, "pck", 4)) {
                debug("Godot executable '%s'", path);
                memcpy(exesections, sections, sizeof(exesections));
                exe = fileopen(path, &size);
                return 0;
            }
    return 1;
}

/**
 * Callback function for find() in godot_load()
 */
int godot_loadcb(char *path)
{
    uint8_t tmp[32], *buf, *buf2, *dbuf;
    uint64_t len, size = 0, subtotal = 0;
    void *f = fileopen(path, &size);
    int i;

    if(f) {
        fileread(f, tmp, sizeof(tmp));
        if(size > 32 && !memcmp(tmp, GDPCMAGIC, 4)) {
            getfileindex(f, 0, size);
        } else
        if(elfpe_getsections(path, 1))
            for(i = 0; sections[i].name[0]; i++)
                if(!memcmp(sections[i].name, "pck", 4)) {
                    getfileindex(f, sections[i].offs, sections[i].size);
                    break;
                }
        for(i = 0; !quit && i < numidx; i++) {
            if((buf = dbuf = (uint8_t*)malloc(idx[i].size + 16))) {
                fileseek(f, idx[i].offs);
                len = fileread(f, buf, idx[i].size);
                if(idx[i].encrypted) len = decrypt_file(buf, len);


                if(!memcmp(buf, "GDST", 4) && len > 0x20) { dbuf += 0x20; len -= 0x20; } else
                if(!memcmp(buf, "GST2", 4) && len > 0x38) { dbuf += 0x38; len -= 0x38; } else
                if(!memcmp(buf, "RSRC", 4) && len > 0x20) {
                    dbuf += 0x20; len -= 0x20;
                    for(; len > 32; dbuf++, len--) {
                        if((!memcmp(dbuf, "RIFF", 4) && !memcmp(dbuf + 8, "WAVE", 4)) ||
                           !memcmp(dbuf, "OggS", 4) || !memcmp(dbuf, "ID3\4", 4) ||
                           (!dbuf[0] && !memcmp(dbuf + 4, "ftyp", 4))) { len -= 4; break; }
                    }
                }
                if(len > 32 && (buf2 = (uint8_t*)malloc(len))) {
                    memcpy(buf2, dbuf, len);
                    if(!fileasset(-1, idx[i].name, buf2, len)) free(buf2);
                }
                free(buf);
            } else {
                engine_features |= EF_NOMEM;
                debug("unable to allocate memory for '%s' (%lu bytes)", idx[i].name, idx[i].size);
            }
            subtotal += idx[i].size;
            status(LOADING, current + subtotal, total);
        }
        fileclose(f);
    }
    current += size;
    status(LOADING, current, total);
    return 1;
}

/**
 * Load all assets from a Godot game
 */
void godot_load(char *path, int isdir)
{
    char tmp[PATH_MAX + FILENAME_MAX + 1];
    char *fn;

    (void)isdir;
    if(!path || !*path) return;

    status(LOADING, 0, 100);
    fn = strrchr(path, '/');
    if(!fn) fn = strrchr(path, '\\');
    if(fn) strncpy(engine_game, fn + 1, sizeof(engine_game) - 1);
    else strcpy(engine_game, "?");

    haskey = 0; memset(godot_key, 0, sizeof(godot_key));
    if(engine_keylen) {
        if(engine_keylen == sizeof(godot_key)) { haskey = 1; memcpy(godot_key, engine_key, sizeof(godot_key)); }
        else msgbox(0, BADKEYLEN, NULL);
    }
    if(haskey)
        debug("encryption key "
            "%02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X "
            "%02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X %02X",
            godot_key[0], godot_key[1], godot_key[2], godot_key[3],
            godot_key[4], godot_key[5], godot_key[6], godot_key[7],
            godot_key[8], godot_key[9], godot_key[10], godot_key[11],
            godot_key[12], godot_key[13], godot_key[14], godot_key[15],
            godot_key[16], godot_key[17], godot_key[18], godot_key[19],
            godot_key[20], godot_key[21], godot_key[22], godot_key[23],
            godot_key[24], godot_key[25], godot_key[26], godot_key[27],
            godot_key[28], godot_key[29], godot_key[30], godot_key[31]);
    else {
        exe = NULL; memset(exesections, 0, sizeof(exesections));
        find(path, 0, godot_findexecb);
    }
    strcpy(tmp, path);
    pathlen = strlen(tmp);
    current = 0; total = du(tmp);
    find(tmp, 0, godot_loadcb);

    engine_features |= EF_DECRYPT;
}

/**
 * Free any additional resources that godot_load() allocated
 */
void godot_unload(void)
{
    freefileindex();
    if(exe) { fileclose(exe); exe = NULL; }
}

/**
 * Callback function for find() in godot_decrypt()
 */
int godot_unpackcb(char *path)
{
    char *fn, *s;
    uint8_t tmp[32], *buf, *dbuf;
    uint64_t len, size = 0, subtotal = 0;
    void *f = fileopen(path, &size);
    int i;

    if(f) {
        fileread(f, tmp, sizeof(tmp));
        if(size > 32 && !memcmp(tmp, GDPCMAGIC, 4)) {
            getfileindex(f, 0, size);
        } else
        if(elfpe_getsections(path, 0))
            for(i = 0; sections[i].name[0]; i++)
                if(!memcmp(sections[i].name, "pck", 4)) {
                    getfileindex(f, sections[i].offs, sections[i].size);
                    break;
                }
        for(fn = path + strlen(path); fn > path + 1 && *fn != '/' && *fn != '\\' && *fn != '.'; fn--);
        if(*fn != '.') *fn++ = '~';
        *fn = 0; filemkdir(path);
        debug("unpacking to '%s'", path);
        *fn++ = '/'; *fn = 0;
        for(i = 0; !quit && i < numidx; i++) {
            if((buf = dbuf = (uint8_t*)malloc(idx[i].size + 16))) {
                fileseek(f, idx[i].offs);
                len = fileread(f, buf, idx[i].size);
                if(idx[i].encrypted) len = decrypt_file(buf, len);
                strcpy(fn, idx[i].name);
                for(s = fn; *s; s++)
                    if(*s == '/') { *s = 0; filemkdir(path); *s = '/'; }
                /* save the original file... */
                putfile(path, buf, len, NULL, 0, NULL, 0);
                if(!(s = strrchr(fn, '.'))) s = fn + strlen(fn);
                /* ...and another one without that stupid Godot header as well */
                if(!memcmp(buf, "GDST", 4) && len > 0x20) { dbuf += 0x20; len -= 0x20; goto img; } else
                if(!memcmp(buf, "GST2", 4) && len > 0x38) {
                    dbuf += 0x38; len -= 0x38;
img:                /* only write out if we are sure about the format */
                    if(!memcmp(dbuf, "RIFF", 4) && !memcmp(dbuf + 8, "WEBP", 4)) { strcpy(s, ".webp"); } else
                    if(!memcmp(dbuf, "\x89PNG", 4)) { strcpy(s, ".png"); } else
                    if(!memcmp(dbuf, "GIF8", 4)) { strcpy(s, ".gif"); } else
                    if(dbuf[0] == 0xff && dbuf[1] == 0xd8 && !memcmp(dbuf + 6, "JFIF", 4)) { strcpy(s, ".jpg"); } else { len = 0; }
                    if(len > 4)
                        putfile(path, dbuf, len, NULL, 0, NULL, 0);
                } else
                if(!memcmp(buf, "RSRC", 4) && len > 0x20) {
                    dbuf += 0x20; len -= 0x20;
                    for(; len > 8; dbuf++, len--) {
                        if(!memcmp(dbuf, "RIFF", 4) && !memcmp(dbuf + 8, "WAVE", 4)) { strcpy(s, ".wav"); len -= 4; break; } else
                        if(!memcmp(dbuf, "OggS", 4)) { strcpy(s, ".ogg"); len -= 4; break; } else
                        if(!memcmp(dbuf, "ID3\4", 4)) { strcpy(s, ".mp3"); len -= 4; break; } else
                        if(!dbuf[0] && !memcmp(dbuf + 4, "ftyp", 4)) { strcpy(s, ".m4a"); len -= 4; break; }
                    }
                    if(len > 4)
                        putfile(path, dbuf, len, NULL, 0, NULL, 0);
                }
                free(buf);
            } else {
                engine_features |= EF_NOMEM;
                debug("unable to allocate memory for '%s' (%lu bytes)", idx[i].name, idx[i].size);
            }
            subtotal += idx[i].size;
            status(SAVING, current + subtotal, total);
        }
        fileclose(f);
    }
    current += size;
    status(SAVING, current, total);
    return 1;
}

/**
 * Decrypt and unpack all files in a Godot game
 */
void godot_decrypt(char *path, int isdir)
{
    char tmp[PATH_MAX + FILENAME_MAX + 1];

    (void)isdir;
    if(!path || !*path) return;

    status(SAVING, 0, 100);
    strcpy(tmp, path);
    current = 0; total = du(tmp);
    find(tmp, 1, godot_unpackcb);
}

/**
 * Draw map tiles to off-screen pixel buffer
 */
void godot_drawmap(map_t *map)
{
    (void)map;
}

/**
 * Save translated text
 */
void godot_translate(char *path)
{
    msgbox(0, UNIMPLEMENTED, path);
}

ENGINE_MODULE(godot, "Godot PCK Archive")
